`define CCM_IDLE 0
`define CCM_X1 1
`define CCM_X2 2
`define CCM_S0 3
`define CCM_Sn 4
`define CCM_Xm 5
`define AES_ENC 6

`define CIPHER_VALID_CNT 11
module aes_ccm(
    input [103:0] nonce_i,
    input [127:0] key_i,
    input [7:0] aad_i,
    input [7:0] payload_length_i,
    input ccm_start_i,
    input [7:0] byte_i,
    input [7:0] byte_cnt_i,
    input byte_update_i,
    output reg byte_out_valid_o,
    output reg [7:0] byte_o,
    output reg [31:0] mic_o,
    input [127:0] aes_data_i,
    output [127:0] aes_rslt_o,
    input aes_start_i,
    output aes_enc_done_o,
    input enc_i,
    input clk,
    input rst_n
);

reg ld;
wire done;
reg [127:0] text_in;
reg [127:0] text_out;
reg [127:0] text_in_r;
wire [127:0] text_out_r;

always @(*) begin
    text_in_r = {
        text_in[7:0],
        text_in[15:8],
        text_in[23:16],
        text_in[31:24],
        text_in[39:32],
        text_in[47:40],
        text_in[55:48],
        text_in[63:56],
        text_in[71:64],
        text_in[79:72],
        text_in[87:80],
        text_in[95:88],
        text_in[103:96],
        text_in[111:104],
        text_in[119:112],
        text_in[127:120]
    };
    text_out = {
        text_out_r[7:0],
        text_out_r[15:8],
        text_out_r[23:16],
        text_out_r[31:24],
        text_out_r[39:32],
        text_out_r[47:40],
        text_out_r[55:48],
        text_out_r[63:56],
        text_out_r[71:64],
        text_out_r[79:72],
        text_out_r[87:80],
        text_out_r[95:88],
        text_out_r[103:96],
        text_out_r[111:104],
        text_out_r[119:112],
        text_out_r[127:120]
    };
end

aes_cipher_top aes_inst(
    .clk(clk),
    .rst(rst_n),
    .ld(ld),
    .done(done),
    .key(key_i),
    .text_in(text_in_r),
    .text_out(text_out_r)
);

reg [2:0] c_st;
reg [2:0] n_st;
always @(posedge clk,negedge rst_n) begin
    if(!rst_n) c_st <= `CCM_IDLE;
    else c_st <= n_st;
end

always @(*) begin
    n_st = c_st;
    case (c_st)
        `CCM_IDLE: if(ccm_start_i) n_st = `CCM_X1;else if(aes_start_i) n_st = `AES_ENC;
        `CCM_X1:   if(done) n_st = `CCM_X2;
        `CCM_X2:   if(done) n_st = `CCM_S0;
        `CCM_S0:   if(done) n_st = `CCM_Sn;
        `CCM_Sn:   if(done) n_st = `CCM_Xm;
        `CCM_Xm:   if(done) n_st = byte_cnt_i == payload_length_i - 1 ? `CCM_IDLE : `CCM_Sn;
        `AES_ENC:  if(done) n_st = `CCM_IDLE;
    endcase
end

reg ld_xm;
always @(posedge clk,negedge rst_n) begin
    if(!rst_n) ld_xm <= 0;
    else if(byte_out_valid_o && (byte_cnt_i%16 == 15||byte_cnt_i == payload_length_i - 1)) ld_xm <= 1;
    else ld_xm <= 0;
end

wire cipher_valid;
assign aes_rslt_o = text_out;
assign aes_enc_done_o = c_st == `AES_ENC && done;
always @(*) begin
    case (c_st)
        `CCM_IDLE: ld = ccm_start_i || aes_start_i;
        `CCM_X1: ld = done;
        `CCM_X2: ld = done;
        `CCM_Sn: ld = byte_update_i && (byte_cnt_i % 16 == 0);
        `CCM_Xm: ld = n_st != `CCM_IDLE ? ld_xm: 0;
        default: ld = 0;
    endcase
end

reg cnt_en;
reg [3:0] cnt;
assign cipher_valid = cnt == `CIPHER_VALID_CNT;

always @(posedge clk,negedge rst_n) begin
    if(!rst_n) byte_out_valid_o <= 0;
    else if(cipher_valid) byte_out_valid_o <= 1;
    else byte_out_valid_o <= 0;
end 
always @(posedge clk,negedge rst_n) begin
    if(!rst_n) cnt_en <= 0;
    else if(byte_update_i) cnt_en <= 1;
    else if(cipher_valid) cnt_en <= 0;
end

always @(posedge clk,negedge rst_n) begin
    if(!rst_n) cnt <= 0;
    else if(cnt_en) cnt <= cnt + 1;
    else cnt <= 0;
end

always @(posedge clk,negedge rst_n) begin
    if(!rst_n) byte_o <= 0;
    else if(cipher_valid) begin
        case (byte_cnt_i%16)
            0:      byte_o <= byte_i ^ text_out[7:0];
            1:      byte_o <= byte_i ^ text_out[15:8];
            2:      byte_o <= byte_i ^ text_out[23:16];
            3:      byte_o <= byte_i ^ text_out[31:24];
            4:      byte_o <= byte_i ^ text_out[39:32];
            5:      byte_o <= byte_i ^ text_out[47:40];
            6:      byte_o <= byte_i ^ text_out[55:48];
            7:      byte_o <= byte_i ^ text_out[63:56];
            8:      byte_o <= byte_i ^ text_out[71:64];
            9:      byte_o <= byte_i ^ text_out[79:72];
            10:     byte_o <= byte_i ^ text_out[87:80];
            11:     byte_o <= byte_i ^ text_out[95:88];
            12:     byte_o <= byte_i ^ text_out[103:96];
            13:     byte_o <= byte_i ^ text_out[111:104];
            14:     byte_o <= byte_i ^ text_out[119:112];
            15:     byte_o <= byte_i ^ text_out[127:120];
        endcase 
    end
end

always @(posedge clk,negedge rst_n) begin
    if(!rst_n) mic_o <= 0;
    else if(c_st==`CCM_S0 && done) mic_o <= text_out_r[127:96];
    else if(c_st==`CCM_Xm && n_st == `CCM_IDLE) mic_o <= mic_o ^ text_out_r[127:96];
end

reg [127:0] msg_blk;
wire [7:0] msg_byte;
assign msg_byte = enc_i ? byte_i : byte_o;
always @(posedge clk,negedge rst_n) begin
    if(!rst_n) msg_blk <= 0;
    else if(enc_i && byte_update_i || (!enc_i && byte_out_valid_o))begin
        case (byte_cnt_i%16)
            0:      msg_blk <= {msg_byte};
            1:      msg_blk[127:8] <= {msg_byte};
            2:      msg_blk[127:16] <= {msg_byte};
            3:      msg_blk[127:24] <= {msg_byte};
            4:      msg_blk[127:32] <= {msg_byte};
            5:      msg_blk[127:40] <= {msg_byte};
            6:      msg_blk[127:48] <= {msg_byte};
            7:      msg_blk[127:56] <= {msg_byte};
            8:      msg_blk[127:64] <= {msg_byte};
            9:      msg_blk[127:72] <= {msg_byte};
            10:     msg_blk[127:80] <= {msg_byte};
            11:     msg_blk[127:88] <= {msg_byte};
            12:     msg_blk[127:96] <= {msg_byte};
            13:     msg_blk[127:104] <= {msg_byte};
            14:     msg_blk[127:112] <= {msg_byte};
            15:     msg_blk[127:120] <= {msg_byte};
        endcase
    end
end

reg [127:0] x_blk;
always @(posedge clk,negedge rst_n) begin
    if(!rst_n) x_blk <= 0;
    else if((c_st == `CCM_X2 || c_st == `CCM_Xm) && done) x_blk <= text_out;
end

always @(*) begin
    case (n_st)
        `CCM_X1: text_in = {payload_length_i,8'h00,nonce_i,8'h49};
        `CCM_X2: text_in = text_out ^ {{13{8'h00}},aad_i,16'h0100};
        `CCM_S0: text_in = {8'h00,8'h00,nonce_i,8'h01};
        `CCM_Sn: text_in = {byte_cnt_i/16+1,8'h00,nonce_i,8'h01};
        `CCM_Xm: text_in = x_blk ^ msg_blk;
        `AES_ENC: text_in = aes_data_i;
        default: text_in = 0;
    endcase
end

endmodule